import json , os , uuid , aiohttp , aiofiles
import time
import traceback
from utils.loggingConfig import get_logger
from service.stepFunService import stepFunService
from fastapi import APIRouter , Request , File, UploadFile , Request
from fastapi.responses import StreamingResponse

fapp = None
router = APIRouter()

root_path = os.path.abspath(__file__)
for i in range(2):
    root_path = os.path.dirname(root_path)

UPLOAD_DIR = os.path.join(root_path , "uploads")
if os.path.exists(UPLOAD_DIR) == False:
    os.makedirs(UPLOAD_DIR)

logger = get_logger('')

# 支持的音频格式
ALLOWED_EXTENSIONS =  {".jpg", ".png", ".jpeg", ".bmp", ".tiff", ".tif", ".webp", ".gif"}

def init():
    global ALLOWED_EXTENSIONS
    ALLOWED_EXTENSIONS = os.getenv('ALLOWED_EXTENSIONS' )
    if ALLOWED_EXTENSIONS :
        ALLOWED_EXTENSIONS = set(ALLOWED_EXTENSIONS.replace(" ","").split(","))
    else:
        ALLOWED_EXTENSIONS = {".jpg", ".png", ".jpeg", ".bmp", ".tiff", ".tif", ".webp", ".gif"}

init()


@router.post("/v1/orc/stepFun", description="体验 orc 模型")
async def stepFun( file: UploadFile = File(...)):
    print("########################################################### stepFun")
    logger.info("######### stepFun")
    req_id = uuid.uuid1()

    # 检查文件扩展名
    print(f"########### file {file}")
    print(f"########### ALLOWED_EXTENSIONS {ALLOWED_EXTENSIONS}")

    file_ext = file.filename.lower()
    if file_ext in ALLOWED_EXTENSIONS:
        # 保存文件
        filename = f'upload-{req_id}-{file.filename}'
        filepath = os.path.join(root_path , f'{UPLOAD_DIR}/{filename}' )
        try:
            with open(filepath, "wb") as f:
                content = await file.read()
                f.write(content)

            ocr_msg = await stepFunService(filepath)
            return {"success": True, "code": 200, 'ocr_msg': ocr_msg }
        except Exception as e:
            error_msg = traceback.format_exc()
            logger.error(f"req_id:{req_id} error_msg:{error_msg}")
            return {"success":False , "code":500 , "error_msg":f"文件上传失败: {error_msg}"}
        finally:
            if os.path.exists(filepath) :
                os.remove(filepath)


async def download_image_async(url: str, save_path):
    """
    异步下载图片

    Args:
        url: 图片URL
        save_path: 保存路径，默认为 error.png
    """
    outputdir = os.path.dirname(save_path)
    if os.path.exists(outputdir) == False:
        os.makedirs(outputdir)
    try:
        async with aiohttp.ClientSession() as session:
            print(f"开始下载: {url}")

            async with session.get(url) as response:
                # 检查响应状态
                if response.status == 200:
                    # 获取文件大小（如果可用）
                    content_length = response.headers.get('Content-Length')
                    if content_length:
                        print(f"文件大小: {int(content_length) / 1024:.2f} KB")

                    # 异步写入文件
                    async with aiofiles.open(save_path, 'wb') as f:
                        async for chunk in response.content.iter_chunked(8192):
                            await f.write(chunk)

                    print(f"下载完成: {save_path}")
                    return True
                else:
                    print(f"下载失败: HTTP {response.status}")
                    return False

    except aiohttp.ClientError as e:
        print(f"网络错误: {e}")
        return False
    except Exception as e:
        print(f"发生错误: {e}")
        return False

@router.post("/v1/chat/completions", description="体验 orc 模型")
async def stepFunApi(request:Request):
    body = await request.json()
    messages = body.get("messages")
    reqId = str(uuid.uuid1()).replace("-","")
    if messages and len(messages) > 0 :
        inputParam = messages[0]
        print("########### messages" , messages)
        try:
            image_url = inputParam["content"][1]["image_url"]["url"]
            if '?' in image_url :
                image_url = image_url.split('?')[0]
        except Exception as e :
            image_url= ""

        filename = image_url.split("/")[-1]
        file_ext = filename.split(".")[-1]
        file_ext = f'.{file_ext}'
        print(f"####### image_url -> {image_url}")
        print(f"####### filename -> {filename}")
        print(f"####### file_ext -> {file_ext}")
        # 确认是否可以传
        if file_ext in ALLOWED_EXTENSIONS:
            filepath = ""
            try:
                filename = image_url.split("/")[-1]
                filename = f'upload-{reqId}-{filename}'
                filepath = os.path.join(root_path, f'{UPLOAD_DIR}/{filename}')
                t1 = time.time()
                await download_image_async(image_url, filepath)
                t2 = time.time()
                logger.info(f'###### modelExperienceController.stepFunApi.downfile_spend_time -> {str(t2-t1)}')
                ocr_msg = await stepFunService(filepath)
                if ocr_msg is None:
                    ocr_msg = ''
                return {"id": f"chatcmpl-{reqId}", "object": "chat.completion", "created": 1763459888,
                        "model": "StepFun/GOT-OCR-2.0-hf",
                        "choices": [{"index": 0, "message": {"role": "assistant", "content": ocr_msg, "refusal": None,
                                                             "annotations": None,
                                                             "audio": None, "function_call": None, "tool_calls": [],
                                                             "reasoning_content": None}, "logprobs": None,
                                     "finish_reason": "length", "stop_reason": None, "token_ids": None}],
                        "service_tier": None, "system_fingerprint": None,
                        "usage": {"prompt_tokens": 1509, "total_tokens": 2048, "completion_tokens": 539,
                                  "prompt_tokens_details": None}, "prompt_logprobs": None,
                        "prompt_token_ids": None, "kv_transfer_params": None}
            except Exception as e :
                errormsg = traceback.format_exc()
                logger.error(f'errormsg->{errormsg}')
            finally:
                if os.path.exists(filepath) :
                    os.remove(filepath)
                    pass

    return {"id": "chatcmpl", "object": "chat.completion", "created": 1763459888, "model": "StepFun/GOT-OCR-2.0-hf", "choices": [{"index": 0, "message": {"role": "assistant", "content": "", "refusal": None,
            "annotations": None, "audio": None, "function_call": None, "tool_calls": [], "reasoning_content": None}, "logprobs": None, "finish_reason": "length", "stop_reason": None, "token_ids": None}],
            "service_tier": None, "system_fingerprint": None, "usage": {"prompt_tokens": 1509, "total_tokens": 2048, "completion_tokens": 539, "prompt_tokens_details": None}, "prompt_logprobs": None,
            "prompt_token_ids": None, "kv_transfer_params": None}
